from numpy import *

def laplace_central_diff_approx(i,j):
  return [(1, i-1, j),
          (1, i+1, j),
	  (1, i, j-1),
	  (1, i, j+1),
	  (-4, i, j)]

def constant_fn(c):
  return lambda *args: c

zero_fn = constant_fn(0)

def q(x):
  x2 = x - 0.5
  return (-4 * x2 * x2) + 1

def build_boundary(top,right,bottom,left):
  def boundary(x,y):
    if y == 1: z = top(x)
    elif x == 1: z = right(y)
    elif y == 0: z = bottom(x)
    elif x == 0: z = left(y)
    else: raise Exception("boundary function called not on a boundary")
    return z
  return boundary

one_on_a_side = build_boundary(constant_fn(1),zero_fn,zero_fn,zero_fn)

one_on_two_sides = build_boundary(constant_fn(1), zero_fn, constant_fn(1), zero_fn)

quad_on_two_sides = build_boundary(q, zero_fn, q, zero_fn)

def solve(approx, g, f, N):
  def on_boundary(i,j):
    return (i==0) or (j==0) or (i==N) or (j==N)

  def flatten(i,j): 
    return (i-1)*(N-1) + (j - 1)

  # useful constants
  fN = float(N)
  h = 1.0 / N
  dim = (N-1)*(N-1) + 1
  one_index = dim - 1

  # build matrix A
  A = zeros( (dim, dim) )
  for i in range(1,N):
    for j in range(1,N):
      curr_row = flatten(i,j)
      for (coeff,x,y) in approx(i,j):
	if on_boundary(x,y):
	  A[curr_row, one_index] += (coeff * g(x/fN,y/fN))
	else:
	  A[curr_row, flatten(x,y)] += coeff
  A[one_index, one_index] = 1
  print "A =\n %s"%A

  # build vector f
  f_vec = zeros(dim)
  for i in range(1,N):
    for j in range(1,N):
      f_vec[flatten(i,j)] = f(i,j)
  f_vec[one_index] = 1
  print "f vector = %s"%f_vec

  # calculate vector u
  A_inv = linalg.inv(A)
  u_vec = dot(A_inv, f_vec)
  print "u vector = %s"%u_vec

  # make function u(i,j)
  def uf(i,j):
    if on_boundary(i,j):
      return g(i/fN,j/fN)
    else:
      return u_vec[flatten(i,j)]

  um = zeros((N+1,N+1))
  for i in range(0,N+1):
    for j in range(0,N+1):
      um[i,j] = uf(i,j)
  return um

